# 一方库

# 官方库

# 三方库
from cProfile import label
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset



class Dataset_siamese(Dataset):
    def __init__(self,dataset_one):
        super(Dataset_siamese, self).__init__()
        self.dataset_one = dataset_one
    def __len__(self):
        return self.dataset_one.__len__()
    def __getitem__(self,index):
        (img,shape,Y,location) = self.dataset_one.__getitem__(index)
        (img_2,shape_2,Y_2,location_2) = None,None,None,None
        # 是否相同(1:1)
        is_same = random.randint(0,1)
        while True:
            # 随机读取一条记录（_2）
            label_target = random.randint(0,self.dataset_one.indexGetter.class_number-1)
            index = self.dataset_one.indexGetter.get_index_random(label_target)
            (path_img_2,Y_2,location_2) = self.dataset_one.__get_path__(index)
            if (Y==Y_2)==is_same: 
                (img_2,shape_2) = self.dataset_one.__load_by_path__(path_img_2)
                is_same  = torch.LongTensor([is_same]) # shape=(1)
                return ((img,shape,Y,location),(img_2,shape_2,Y_2,location_2))

class Model_siamese(nn.Module):
    """孪生姊妹网络"""
    def __init__(self,model_one):
        super(Model_siamese, self).__init__()
        self.model_one = model_one
        # print(self.model_one)
    def forward(self,input,shape,input2,shape2):
        return (self.model_one(input,shape),self.model_one(input2,shape2))

class Loss_contrastive(torch.nn.Module):
    """ContrastiveLoss"""
    def __init__(self, margin=2.0):
        super(Loss_contrastive, self).__init__()
        self.margin = margin

    def forward(self, output1, output2, is_same): # is_same.shape=[-1]
        euclidean_distance = F.pairwise_distance(output1, output2)
        loss_contrastive = torch.mean((is_same) * torch.pow(euclidean_distance, 2) +
                                      (1-is_same) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))
        return loss_contrastive*0.5

class Loss_siamese(torch.nn.Module):
    """孪生姊妹网络loss"""
    def __init__(self,loss_one,margin):
        super(Loss_siamese, self).__init__()
        self.loss_one = loss_one # 单网络loss
        self.loss_contrastive = Loss_contrastive(margin)  # ContrastiveLoss
    def forward(self,embedding  ,Y_hat  ,Y  ,location_hat  ,location  , 
                     embedding_2,Y_hat_2,Y_2,location_hat_2,location_2,):
        loss_one   = self.loss_one(Y_hat  ,Y  ,location_hat  ,location) 
        loss_one_2 = self.loss_one(Y_hat_2,Y_2,location_hat_2,location_2) 

        loss_one_mean = (loss_one + loss_one_2)/2
        loss_contrastive = self.loss_contrastive(embedding,embedding_2,(Y==Y_2).reshape(-1).int())
        
        loss_siamese = loss_one_mean+loss_contrastive
        
        return loss_siamese

